# -*- coding: utf-8 -*-
# Author: Sun Jiawei
# E-mail: sunjiawei@tbea.com
#
import sys
import warnings

import pywt
import matplotlib.pyplot as plt

from visualization import DataSet

warnings.filterwarnings('ignore')


def plot_dwt(data, level):
    wavelet_name = 'db1'
    # ca1, cd1 = pywt.wavedec(data, wavelet_name, level=1)
    a = data
    w = pywt.Wavelet(wavelet_name)
    mode = pywt.Modes.smooth
    ca = []
    cd = []
    for _ in range(level):
        a, d = pywt.dwt(a, w, mode)
        ca.append(a)
        cd.append(d)

    rec_a = []
    rec_d = []
    for i, coeff in enumerate(ca):
        coeff_list = [coeff, None] + [None] * i
        rec_a.append(pywt.waverec(coeff_list, w))

    for i, coeff in enumerate(cd):
        coeff_list = [None, coeff] + [None] * i
        rec_d.append(pywt.waverec(coeff_list, w))

    fig = plt.figure()
    ax_main = fig.add_subplot(level+1, 1, 1)
    ax_main.set_title('原始电价序列(2015全年1:00电价)')
    ax_main.plot(data)
    ax_main.set_xlim(0, len(data)-1)

    for i, y in enumerate(rec_a):
        ax = fig.add_subplot(len(rec_a) + 1, 2, 3 + i * 2)
        ax.plot(y, 'r')
        ax.set_xlim(0, len(y) - 1)
        ax.set_ylabel("A%d" % (i + 1))

    for i, y in enumerate(rec_d):
        ax = fig.add_subplot(len(rec_d) + 1, 2, 4 + i * 2)
        ax.plot(y, 'g')
        ax.set_xlim(0, len(y) - 1)
        ax.set_ylabel("D%d" % (i + 1))


if __name__ == '__main__':
    level = int(sys.argv[1])
    assert level >= 1
    file_path = './data/opsd-time_series-2020-10-06/time_series_15min_singleindex.csv'
    ds = DataSet(file_path)
    h1, m1 = '1', '0'
    df1 = ds.get_panel_data('2015', '1', '2016', '1', h1, m1)
    x = df1.index
    y1 = df1['AT_price_day_ahead']
    plot_dwt(y1.values, level)
    plt.show()

